import sqlite3
import os
from datetime import datetime

from config import DATABASE

# 数据库版本定义
DB_VERSIONS = [
    {
        'version': 1,
        'description': '初始数据库结构',
        'sql': [
            '''CREATE TABLE IF NOT EXISTS users (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                username TEXT UNIQUE NOT NULL,
                email TEXT UNIQUE NOT NULL,
                password_hash TEXT NOT NULL,
                role TEXT DEFAULT 'user',
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
            )''',
            '''CREATE TABLE IF NOT EXISTS content (
                id TEXT PRIMARY KEY,
                user_id INTEGER NOT NULL,
                type TEXT NOT NULL,
                content TEXT,
                file_path TEXT,
                url TEXT,
                title TEXT,
                extracted_text TEXT,
                keywords TEXT,
                summary TEXT,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (user_id) REFERENCES users (id)
            )''',
            '''CREATE TABLE IF NOT EXISTS comments (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content_id TEXT NOT NULL,
                user_id INTEGER NOT NULL,
                comment TEXT NOT NULL,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (content_id) REFERENCES content (id),
                FOREIGN KEY (user_id) REFERENCES users (id)
            )'''
        ]
    },
    {
        'version': 2,
        'description': '添加评论回复功能',
        'sql': [
            '''ALTER TABLE comments ADD COLUMN parent_id INTEGER REFERENCES comments(id)'''
        ]
    },
    {
        'version': 3,
        'description': '添加私信功能',
        'sql': [
            '''CREATE TABLE IF NOT EXISTS messages (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                sender_id INTEGER NOT NULL,
                recipient_id INTEGER NOT NULL,
                subject TEXT NOT NULL,
                content TEXT NOT NULL,
                is_read BOOLEAN DEFAULT FALSE,
                read_at TIMESTAMP,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (sender_id) REFERENCES users (id),
                FOREIGN KEY (recipient_id) REFERENCES users (id)
            )'''
        ]
    },
    {
        'version': 4,
        'description': '添加数据库版本控制表',
        'sql': [
            '''CREATE TABLE IF NOT EXISTS db_version (
                version INTEGER PRIMARY KEY,
                applied_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                description TEXT
            )'''
        ]
    },
    {
        'version': 5,
        'description': '添加内容索引以提高搜索性能',
        'sql': [
            '''CREATE INDEX IF NOT EXISTS idx_content_keywords ON content(keywords)''',
            '''CREATE INDEX IF NOT EXISTS idx_content_type ON content(type)''',
            '''CREATE INDEX IF NOT EXISTS idx_content_user ON content(user_id)''',
            '''CREATE INDEX IF NOT EXISTS idx_comments_content ON comments(content_id)''',
            '''CREATE INDEX IF NOT EXISTS idx_comments_user ON comments(user_id)'''
        ]
    },
    {
        'version': 6,
        'description': '添加Markdown支持',
        'sql': [
            '''ALTER TABLE content ADD COLUMN html_content TEXT''',
            '''ALTER TABLE content ADD COLUMN raw_content TEXT'''
        ]
    },
    {
        'version': 7,
        'description': '添加批注功能',
        'sql': [
            '''CREATE TABLE IF NOT EXISTS annotations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                content_id TEXT NOT NULL,
                user_id INTEGER NOT NULL,
                type TEXT NOT NULL,
                anchor_json TEXT NOT NULL,
                note TEXT NOT NULL,
                is_private BOOLEAN DEFAULT FALSE,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (content_id) REFERENCES content (id),
                FOREIGN KEY (user_id) REFERENCES users (id)
            )''',
            '''CREATE INDEX IF NOT EXISTS idx_annotations_content ON annotations(content_id)''',
            '''CREATE INDEX IF NOT EXISTS idx_annotations_user ON annotations(user_id)'''
        ]
    }
]

def get_db_connection():
    """获取数据库连接"""
    conn = sqlite3.connect(DATABASE)
    conn.row_factory = sqlite3.Row
    return conn

def get_current_db_version():
    """获取当前数据库版本"""
    try:
        conn = get_db_connection()
        
        # 检查版本表是否存在
        cursor = conn.execute("""
            SELECT name FROM sqlite_master 
            WHERE type='table' AND name='db_version'
        """)
        
        if not cursor.fetchone():
            conn.close()
            return 0
        
        # 获取最新版本
        cursor = conn.execute("SELECT MAX(version) as version FROM db_version")
        result = cursor.fetchone()
        conn.close()
        
        return result['version'] if result['version'] else 0
        
    except Exception as e:
        print(f"获取数据库版本失败: {e}")
        return 0

def apply_migration(version_info):
    """应用单个迁移"""
    try:
        conn = get_db_connection()
        
        print(f"应用迁移 v{version_info['version']}: {version_info['description']}")
        
        for sql in version_info['sql']:
            try:
                conn.execute(sql)
                print(f"  执行SQL: {sql[:50]}...")
            except Exception as e:
                print(f"  SQL执行失败: {e}")
                print(f"  SQL: {sql}")
                # 对于某些ALTER TABLE操作，如果列已存在会失败，这是正常的
                if "duplicate column name" not in str(e).lower():
                    raise e
        
        # 只有在db_version表存在时才记录版本
        try:
            conn.execute("""
                INSERT OR REPLACE INTO db_version (version, description) 
                VALUES (?, ?)
            """, (version_info['version'], version_info['description']))
        except Exception as e:
            # 如果db_version表不存在（版本1-3），忽略这个错误
            if "no such table: db_version" not in str(e):
                raise e
            print(f"  跳过版本记录（db_version表尚未创建）")
        
        conn.commit()
        conn.close()
        
        print(f"迁移 v{version_info['version']} 完成")
        return True
        
    except Exception as e:
        print(f"应用迁移 v{version_info['version']} 失败: {e}")
        return False

def init_database():
    """初始化数据库并应用所有迁移"""
    try:
        # 确保数据目录存在
        os.makedirs('data', exist_ok=True)
        
        print("开始数据库初始化...")
        
        current_version = get_current_db_version()
        latest_version = DB_VERSIONS[-1]['version']
        
        print(f"当前版本: {current_version}, 最新版本: {latest_version}")
        
        # 确保版本号是整数类型
        current_version = int(current_version) if current_version is not None else 0
        latest_version = int(latest_version)
        
        if current_version >= latest_version:
            print("数据库已是最新版本")
            return verify_database_integrity()
        
        # 应用需要的迁移
        for version_info in DB_VERSIONS:
            if int(version_info['version']) > current_version:
                if not apply_migration(version_info):
                    print(f"迁移失败，停止在版本 {version_info['version']}")
                    return False
        
        print("数据库迁移完成")
        return verify_database_integrity()
        
    except Exception as e:
        print(f"数据库初始化失败: {e}")
        return False

def verify_database_integrity():
    """验证数据库完整性"""
    try:
        conn = get_db_connection()
        
        # 检查必要的表是否存在
        required_tables = ['users', 'content', 'comments', 'messages', 'annotations', 'db_version']
        
        for table in required_tables:
            cursor = conn.execute("""
                SELECT name FROM sqlite_master 
                WHERE type='table' AND name=?
            """, (table,))
            
            if not cursor.fetchone():
                print(f"缺少必要的表: {table}")
                conn.close()
                return False
        
        # 检查关键字段是否存在
        cursor = conn.execute("PRAGMA table_info(content)")
        columns = [row[1] for row in cursor.fetchall()]
        
        required_columns = ['id', 'user_id', 'type', 'content', 'html_content', 'raw_content']
        for column in required_columns:
            if column not in columns:
                print(f"content表缺少字段: {column}")
                conn.close()
                return False
        
        conn.close()
        print("数据库完整性检查通过")
        return True
        
    except Exception as e:
        print(f"数据库完整性检查失败: {e}")
        return False

def repair_database():
    """修复数据库"""
    try:
        print("开始修复数据库...")
        
        # 备份当前数据库
        if os.path.exists(DATABASE):
            backup_path = f"{DATABASE}.backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            import shutil
            shutil.copy2(DATABASE, backup_path)
            print(f"数据库已备份到: {backup_path}")
        
        # 重新初始化数据库
        return init_database()
        
    except Exception as e:
        print(f"数据库修复失败: {e}")
        return False

def reset_database():
    """重置数据库（危险操作）"""
    try:
        if os.path.exists(DATABASE):
            # 备份
            backup_path = f"{DATABASE}.reset_backup.{datetime.now().strftime('%Y%m%d_%H%M%S')}"
            import shutil
            shutil.copy2(DATABASE, backup_path)
            print(f"数据库已备份到: {backup_path}")
            
            # 删除原数据库
            os.remove(DATABASE)
            print("原数据库已删除")
        
        # 重新创建
        return init_database()
        
    except Exception as e:
        print(f"重置数据库失败: {e}")
        return False

if __name__ == '__main__':
    print("数据库迁移工具")
    print("1. 初始化/更新数据库")
    print("2. 检查数据库完整性")
    print("3. 修复数据库")
    print("4. 重置数据库（危险）")
    
    choice = input("请选择操作 (1-4): ").strip()
    
    if choice == '1':
        if init_database():
            print("数据库初始化/更新成功")
        else:
            print("数据库初始化/更新失败")
    elif choice == '2':
        if verify_database_integrity():
            print("数据库完整性检查通过")
        else:
            print("数据库完整性检查失败")
    elif choice == '3':
        if repair_database():
            print("数据库修复成功")
        else:
            print("数据库修复失败")
    elif choice == '4':
        confirm = input("确定要重置数据库吗？这将删除所有数据！(yes/no): ").strip().lower()
        if confirm == 'yes':
            if reset_database():
                print("数据库重置成功")
            else:
                print("数据库重置失败")
        else:
            print("操作已取消")
    else:
        print("无效选择")